package com.bigmodel.baidu;

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpUtil;
import cn.hutool.json.JSONUtil;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class BaiduWenXin {

    private final static String API_KEY = "r4UzD8T06ld7RWo5QfKDeOKb";
    private final static String  SECRET_KEY = "tMawXUmwgSbKkpGAL1HwffsEDKOyPZhO";


    public static String getToken(){
        String url = "https://aip.baidubce.com/oauth/2.0/token";
        Map<String,Object> map  = new HashMap<>();
        map.put("grant_type","client_credentials");
        map.put("client_id",API_KEY);
        map.put("client_secret",SECRET_KEY);
        String response = HttpUtil.get(url, map);

        if(ObjectUtil.isNotEmpty(response)){
            return JSONUtil.parseObj(response).getStr("access_token");
        }
        return null;
    }


    public static String chat(String prompt){

        String token = getToken();
        if(ObjectUtil.isEmpty(token)){
            return "501";
        }

        //String url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions?access_token=" + token;
        String url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant?access_token=" + token;

        Map<String,Object> params = new HashMap<>();
        List<Map<String,String>> messages = new ArrayList<>();
        
        Map<String,String> messageItem = new HashMap<>();
        messageItem.put("role","user");
        messageItem.put("content",prompt);
        messages.add(messageItem);
        params.put("messages",messages);
        params.put("temperature",0.01);

        HttpRequest reuqest = HttpUtil.createPost(url);
        Map<String,String> headers = new HashMap<>();
        headers.put("Content-Type","application/json");
        reuqest.addHeaders(headers);
        reuqest.body(JSONUtil.toJsonStr(params));

        //System.out.println(prompt);
        //System.out.println("---------------------------divider---------------------------");

        String response = reuqest.execute().body();

        if(ObjectUtil.isNotEmpty(response)){
            return JSONUtil.parseObj(response).getStr("result");
        }

        return null;
    }


    public static String chat(List<String> prompt){

        String token = getToken();
        if(ObjectUtil.isEmpty(token)){
            return "501";
        }

       // String url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions?access_token=" + token;
        String url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant?access_token=" + token;

        Map<String,Object> params = new HashMap<>();
        List<Map<String,String>> messages = new ArrayList<>();

        for(int i = 0 ; i< prompt.size(); i++) {
            if(i%2==0) {
                Map<String, String> messageItem = new HashMap<>();
                messageItem.put("role", "user");
                messageItem.put("content", prompt.get(i));
                messages.add(messageItem);
            }else{
                Map<String, String> messageItem = new HashMap<>();
                messageItem.put("role", "assistant");
                messageItem.put("content", prompt.get(i));
                messages.add(messageItem);
            }
        }
        params.put("messages",messages);
//        params.put("temperature",0.1);

        HttpRequest reuqest = HttpUtil.createPost(url);
        Map<String,String> headers = new HashMap<>();
        headers.put("Content-Type","application/json");
        reuqest.addHeaders(headers);
        reuqest.body(JSONUtil.toJsonStr(params));

        //System.out.println(prompt);
        //System.out.println("---------------------------divider---------------------------");

        String response = reuqest.execute().body();

        if(ObjectUtil.isNotEmpty(response)){
            return JSONUtil.parseObj(response).getStr("result");
        }

        return null;
    }



    public static String chatList(List<Map<String,String>> messages){

        String token = getToken();
        if(ObjectUtil.isEmpty(token)){
            return "501";
        }
        // String url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions?access_token=" + token;
        String url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant?access_token=" + token;

        Map<String,Object> params = new HashMap<>();
        params.put("messages",messages);
        params.put("temperature",0.1);

        HttpRequest reuqest = HttpUtil.createPost(url);
        Map<String,String> headers = new HashMap<>();
        headers.put("Content-Type","application/json");
        reuqest.addHeaders(headers);
        reuqest.body(JSONUtil.toJsonStr(params));

        //System.out.println(prompt);
        //System.out.println("---------------------------divider---------------------------");

        String response = reuqest.execute().body();

        if(ObjectUtil.isNotEmpty(response)){
            return JSONUtil.parseObj(response).getStr("result");
        }

        return null;
    }


    public static void main(String[] args) {
        String chat = BaiduWenXin.chat("你好,帮我查下中国10大热门网站");
        System.out.println(chat);
    }



}
